{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Genetic algorithm with the help of functional operators\n",
    "\n",
    "In this notebook, we demonstrate how one can design a genetic algorithm by combining operators implemented within the namespace `evotorch.operators.functional`.\n",
    "\n",
    "## Introduction\n",
    "\n",
    "EvoTorch provides the namespace `evotorch.operators.functional` which contains genetic algorithm operators that can be called directly on PyTorch tensors. These operators are implemented in conformance with the functional programming paradigm, meaning that they do not mutate the tensors they receive as arguments (`*`).\n",
    "\n",
    "A genetic algorithm can be designed by simply calling these genetic algorithm operators in an evolution loop. This way of implementing a genetic algorithm grants the user complete flexibility regarding how and when the operators are to be called, and what extra procedures are to be followed between these operator calls.\n",
    "\n",
    "## Use cases\n",
    "\n",
    "**Batched optimization.**\n",
    "These operators are defined in such a way that, if they receive a population tensor with 3 or more dimensions (instead of 2 dimensions), the extra leftmost dimensions are interpreted as batch dimensions, and the steps of the operators are broadcast to those batch dimensions. This means that they can work not just on a population, but on a batch of populations, in a vectorized manner.\n",
    "\n",
    "This feature could be helpful when one has multiple populations (each initialized around different values and/or using different initialization methods), and one wishes to run an evolutionary search on all these populations efficiently.\n",
    "\n",
    "**Nested optimization.**\n",
    "It could be the case that the optimization problem at hand has an inner (nested) optimization problem that needs to be addressed within its fitness function. In such cases, one has to run an inner evolutionary search while evaluating each solution. This inner evolutionary search could be implemented with the help of these functional operators. Considering that each solution of the outer problem ends up with its own inner optimization problem, this way of tackling the inner problem could result in an efficient and vectorized implementation (vectorization would happen across multiple inner optimization problems induced by multiple solutions of the outer problem; see the use case \"Batched optimization\").\n",
    "\n",
    "---\n",
    "\n",
    "`(*)` It is to be noted, however, that they _do_ mutate the global random state of PyTorch, because of how they use PyTorch functions such as `torch.randn(...)`, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Example\n",
    "\n",
    "We now show how to use the functional operators to design a genetic algorithm to solve the Rastrigin problem. To keep the example simple, we do not consider the use cases of batched/nested optimization.\n",
    "\n",
    "We begin with the necessary imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from evotorch.operators import functional as func_ops\n",
    "from evotorch.decorators import rowwise\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Below, we have the implementations for the fitness functions `rastrigin` and `sphere`.\n",
    "\n",
    "Notice how these fitness functions are decorated via `evotorch.decorators.rowwise`.\n",
    "This decorator allows the user to implement the function with the assumption that its received argument is a single row (i.e. a 1-dimensional tensor). As an additional behavior, if a function decorated via `@rowwise` receives a tensor with 2 or more dimensions, the operations defined within the decorated function are broadcast across the extra leftmost dimensions. In other words, the extra leftmost dimensions are interpreted as batch dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@rowwise\n",
    "def rastrigin(x: torch.Tensor) -> torch.Tensor:\n",
    "    from math import pi\n",
    "    A = 10\n",
    "    [n] = x.shape\n",
    "    return A * n + torch.sum((x ** 2.0) - (A * torch.cos(2 * pi * x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@rowwise\n",
    "def sphere(x: torch.Tensor) -> torch.Tensor:\n",
    "    return torch.linalg.norm(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "In this notebook, the variable `f` points to the fitness function whose value we want to minimize:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#f = sphere\n",
    "f = rastrigin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Various hyperparameters and problem settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "popsize = 1000  # population size\n",
    "solution_length = 1000  # length of a solution\n",
    "\n",
    "# lower and upper bounds for the decision values of the initial population\n",
    "lb = -5.12\n",
    "ub = 5.12\n",
    "\n",
    "tournament_size = 8  # tournament size\n",
    "mutation_stdev = 0.01  # standard deviation for the Gaussian mutation\n",
    "eta = 10.0  # eta value for the simulated binary cross-over\n",
    "\n",
    "num_generations = 1000  # number of generations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the decision values of the initial population\n",
    "parents = (torch.rand(popsize, solution_length, dtype=torch.float32) * (ub - lb)) + lb\n",
    "\n",
    "# Evaluate the initial population\n",
    "parent_evals = f(parents)\n",
    "\n",
    "last_reporting_time = None\n",
    "reporting_interval = 1\n",
    "\n",
    "# Main loop of the population\n",
    "for generation in range(1, 1 + num_generations):\n",
    "\n",
    "    # Given the parent solutions and their evaluation results,\n",
    "    # run a tournament, pick pairs, and apply cross-over for each pair:\n",
    "    candidates = func_ops.simulated_binary_cross_over(\n",
    "        parents,\n",
    "        parent_evals,\n",
    "        eta=eta,\n",
    "        tournament_size=tournament_size,\n",
    "        objective_sense=\"min\",\n",
    "    )\n",
    "\n",
    "    # Instead of simulated binary cross-over, you could use two-point\n",
    "    # cross-over:\n",
    "    # candidates = func_ops.two_point_cross_over(\n",
    "    #     parents,\n",
    "    #     parent_evals,\n",
    "    #     tournament_size=tournament_size,\n",
    "    #     objective_sense=\"min\",\n",
    "    # )\n",
    "\n",
    "    # Apply Gaussian mutation on the newly made candidate solutions\n",
    "    candidates = candidates + (torch.randn_like(candidates) * mutation_stdev)\n",
    "\n",
    "    # On the newly mutated solutions, apply the permutation operator of the CoSyNE algorithm\n",
    "    permuted = func_ops.cosyne_permutation(parents, permute_all=True)\n",
    "    # Add the permutation results onto the new population of candidates\n",
    "    candidates = func_ops.combine(candidates, permuted)\n",
    "\n",
    "    # Evaluate all the candidate solutions\n",
    "    candidate_evals = f(candidates)\n",
    "\n",
    "    # Combine the parent population and the candidate population to form an\n",
    "    # extended population. This time, we combine together with the evaluation results.\n",
    "    extended_population, extended_evals = (\n",
    "        func_ops.combine((parents, parent_evals), (candidates, candidate_evals))\n",
    "    )\n",
    "\n",
    "    # From the extended population, take the best `popsize` number of solutions.\n",
    "    # These taken solutions will server as the parents of the next generation.\n",
    "    parents, parent_evals = (\n",
    "        func_ops.take_best(extended_population, extended_evals, popsize, objective_sense=\"min\")\n",
    "    )\n",
    "\n",
    "    # Report how the evolution is progressing\n",
    "    now = datetime.now()\n",
    "    if (\n",
    "        (last_reporting_time is None)\n",
    "        or (generation == num_generations)\n",
    "        or ((now - last_reporting_time).total_seconds() > reporting_interval)\n",
    "    ):\n",
    "        last_reporting_time = now\n",
    "        print(\"Generation:\", generation, \" Best eval of population:\", parent_evals.min())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "Decision values of the final population:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "parents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "Best solution of the final population:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_best, pop_best_eval = func_ops.take_best(parents, parent_evals, objective_sense=\"min\")\n",
    "\n",
    "print(\"Best solution of the final population:\", pop_best)\n",
    "print(\"Best evaluation result of the final population:\", pop_best_eval)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
